# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: https://docs.scrapy.org/en/latest/topics/item-pipeline.html
import os

# useful for handling different item types with a single interface
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from eastmoney.entities import Rate, Nav, Base
from eastmoney.items import RateItem, NavItem


class EastmoneyPipeline:

    def __init__(self):
        """
            初始化数据库、创建连接
        """

        mysql_host = os.environ.get('MYSQL_HOST', 'localhost')
        mysql_port = os.environ.get('MYSQL_PORT', '3307')
        mysql_user = os.environ.get('MYSQL_USER', 'root')
        mysql_password = os.environ.get('MYSQL_PASSWORD', '123456')
        mysql_db = os.environ.get('MYSQL_DB', 'eastmoney')

        db_url = f"mysql+pymysql://{mysql_user}:{mysql_password}@{mysql_host}:{mysql_port}/{mysql_db}"
        print(f"数据库连接url： {db_url}")
        self.engine = create_engine(db_url)

        Session = sessionmaker(bind=self.engine)
        self.session = Session()

        Base.metadata.drop_all(self.engine)
        Base.metadata.create_all(self.engine)

        self.fund_buffer = []
        self.nav_buffer = []
        self.codeDict={}

    def open_spider(self, spider):
        pass

    def close_spider(self, spider):
        if len(self.fund_buffer) > 0:
            self.session.bulk_save_objects(self.fund_buffer)
            self.session.commit()
            # print(f"insert left fund recodes")
            self.fund_buffer.clear()
        if len(self.nav_buffer) > 0:
            self.session.bulk_save_objects(self.nav_buffer)
            self.session.commit()
            # print(f"insert left nav recodes")
            self.nav_buffer.clear()

        print(f"抓取基金总数：{len(self.codeDict)}")
        self.session.close()

    def batch_insert_nav(self, buffer):
        """
        将缓冲区的数据刷到DB
        :param buffer: 缓冲区数据
        :return:
        """
        self.session.bulk_save_objects(buffer)
        self.session.commit()
        buffer.clear()

    def batch_insert_fund(self, buffer):
        """
        将缓冲区的数据刷到DB
        :param buffer: 缓冲区数据
        :return:
        """
        self.session.bulk_save_objects(buffer)
        self.session.commit()
        buffer.clear()

    def process_item(self, item, spider):
        self.codeDict[item['fund_code']] = item['fund_code']

        if isinstance(item, RateItem):
            self.process_rate_item(item, spider)
        elif isinstance(item, NavItem):
            self.process_nav_item(item, spider)

        if len(self.fund_buffer) >= 300:
            self.batch_insert_fund(self.fund_buffer)
        if len(self.nav_buffer) >= 300:
            self.batch_insert_nav(self.nav_buffer)

        return item

    def process_rate_item(self, item, spider):
        self.fund_buffer.append(Rate(**item))
        return item

    def process_nav_item(self, item, spider):
        self.nav_buffer.append(Nav(**item))
        return item
